package middleware

import (
	"context"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis_rate/v9"
	"golang.org/x/time/rate"
	"tieup/internal/config"
	"tieup/pkg/redis"
)

// 基于内存的速率限制器变量
var (
	mu         sync.Mutex
	ipLimiters = make(map[string]*RateLimiter)
)

// 添加 lastSeen 字段到 RateLimiter 结构
type RateLimiter struct {
	limiter  *rate.Limiter
	lastSeen time.Time
}

// 全局速率限制中间件
func RateLimitMiddleware(cfg config.RateLimitConfig) gin.HandlerFunc {
	if !cfg.Enabled {
		return func(c *gin.Context) { c.Next() }
	}

	if cfg.Storage == "redis" {
		return rateLimitRedis(cfg)
	}
	return rateLimitLocal(cfg)
}

// 基于Redis的分布式速率限制
func rateLimitRedis(cfg config.RateLimitConfig) gin.HandlerFunc {
	limiter := redis_rate.NewLimiter(redis.Client)
	limit := redis_rate.Limit{
		Period: time.Second * time.Duration(cfg.Duration),
		Rate:   cfg.Requests,
		Burst:  cfg.Requests * 2,
	}

	return func(c *gin.Context) {
		ip := c.ClientIP()
		res, err := limiter.Allow(context.Background(), "rate_limit:"+ip, limit)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "Rate limit service unavailable"})
			return
		}

		// 设置速率限制头部信息
		c.Header("X-RateLimit-Limit", strconv.Itoa(cfg.Requests))
		c.Header("X-RateLimit-Remaining", strconv.Itoa(res.Remaining))
		c.Header("X-RateLimit-Reset", strconv.FormatInt(res.ResetAfter.Milliseconds()/1000, 10))

		if res.Allowed == 0 {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error":       "Too many requests",
				"retry_after": res.RetryAfter.Seconds(),
			})
			return
		}

		c.Next()
	}
}

// 基于内存的速率限制
func rateLimitLocal(cfg config.RateLimitConfig) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		mu.Lock()
		limiter, exists := ipLimiters[ip]
		if !exists {
			limiter = &RateLimiter{
				limiter:  rate.NewLimiter(rate.Every(time.Duration(cfg.Duration)*time.Second), cfg.Requests),
				lastSeen: time.Now(),
			}
			ipLimiters[ip] = limiter
		} else {
			// 更新最后访问时间
			limiter.lastSeen = time.Now()
		}
		mu.Unlock()

		if !limiter.limiter.Allow() {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests"})
			return
		}

		c.Next()
	}
}

// 定期清理过期的限流器
func cleanupExpiredLimiters() {
	for {
		time.Sleep(10 * time.Minute)
		mu.Lock()
		for ip, limiter := range ipLimiters {
			if time.Since(limiter.lastSeen) > 30*time.Minute {
				delete(ipLimiters, ip)
			}
		}
		mu.Unlock()
	}
}

// 启动清理goroutine
func init() {
	go cleanupExpiredLimiters()
}
